import sys
from dataclasses import dataclass
from typing import Any, Callable, Optional

from paramiko.client import SSHClient
from paramiko.ssh_exception import NoValidConnectionsError
from zkl_serialization import JsonValue, dump_json_value, parse_json_value

from jobs_agency.skeletons import SkeletonAttachedJob, SkeletonDetachedJob, SkeletonPlatform, StoppedJob
from .config import SSHConfig
from .utils_life import make_check_script, make_life_script, make_monitor_script, make_terminate_script
from .utils_ssh import sftp_upload, ssh_call, ssh_connect, ssh_execute


@dataclass
class SSHPlan:
    command_script: str = ""
    prepare_func: Optional[Callable[[SSHClient], Any]] = None
    release_script: str = ""
    work_dir_path: Optional[str] = None
    logs_dir_path: Optional[str] = None
    job_name: Optional[str] = None


class SSHAttachedJob(SkeletonAttachedJob):
    def __init__(self, ssh_config: SSHConfig):
        super().__init__()
        self._ssh_config = ssh_config
        self._ssh_client = ssh_connect(ssh_config)

    @property
    def ssh_config(self) -> SSHConfig:
        return self._ssh_config

    @property
    def ssh_client(self) -> SSHClient:
        return self._ssh_client

    def _detach(self):
        self._ssh_client = None

    def _terminate(self):
        self._ssh_client.close()
        self._ssh_client = None


class SSHDetachedJob(SkeletonDetachedJob):
    def __init__(self,
        id: str, platform: 'SSHPlatform',
        ssh_client: SSHClient, work_dir_path: str,
    ):
        super().__init__(id, platform)
        self._ssh_client = ssh_client
        self._work_dir_path = work_dir_path

    def _check_active(self) -> bool:
        state_str = ssh_call(self._ssh_client, f"{self._work_dir_path}/check.sh").strip()
        return state_str != "STOPPED"

    def _monitor(self):
        ssh_execute(self._ssh_client, f"{self._work_dir_path}/monitor.sh")

    def _terminate(self):
        ssh_execute(self._ssh_client, f"{self._work_dir_path}/terminate.sh")


class SSHPlatform(SkeletonPlatform[SSHConfig, SSHPlan, SSHAttachedJob, SSHDetachedJob]):
    def _new_attached_job(self, config: SSHConfig) -> SSHAttachedJob:
        return SSHAttachedJob(config)

    def _new_detached_job(self, config: SSHConfig, plan: SSHPlan) -> SSHDetachedJob:
        with self._new_attached_job(config) as attached_job:
            detached_job = self._new_detached_job_from_attached_job(attached_job, plan)
            attached_job.detach()
            return detached_job

    def _new_detached_job_from_attached_job(self, attached_job: SSHAttachedJob, plan: SSHPlan) -> SSHDetachedJob:
        config = attached_job.ssh_config
        client = attached_job.ssh_client

        work_dir_path = plan.work_dir_path or ssh_call(client, "mktemp -d").strip()
        logs_dir_path = plan.logs_dir_path or "."
        job_id = self._store.save_job(None, {"ssh": dump_json_value(config), 'work_dir_path': work_dir_path})

        if plan.prepare_func:
            print("Preparing...", file=sys.stderr)
            plan.prepare_func(client)
            print("", file=sys.stderr)

        print("Setting up workplace...", file=sys.stderr)
        ssh_execute(client, (
            f"set -e \n"
            f"mkdir -p {work_dir_path} \n"
            f"cd {work_dir_path} \n"
            f"mkdir -p {logs_dir_path} \n"))
        with client.open_sftp() as sftp:
            sftp.chdir(work_dir_path)
            sftp_upload(sftp, plan.command_script.encode(), "command.sh")
            sftp_upload(sftp, plan.release_script.encode(), "release.sh")
            sftp_upload(sftp, make_life_script(work_dir_path, logs_dir_path).encode(), "life.sh")
            sftp_upload(sftp, make_check_script(work_dir_path, logs_dir_path).encode(), "check.sh")
            sftp_upload(sftp, make_monitor_script(work_dir_path, logs_dir_path).encode(), "monitor.sh")
            sftp_upload(sftp, make_terminate_script(work_dir_path, logs_dir_path).encode(), "terminate.sh")
        ssh_execute(client, f"chmod +x {work_dir_path}/*.sh")
        print("", file=sys.stderr)

        print("Starting life...", file=sys.stderr)
        ssh_execute(client, (
            f"set -e \n"
            f"cd {work_dir_path} \n"
            f"nohup ./life.sh "
            f" >> {logs_dir_path}/stdout.txt "
            f"2>> {logs_dir_path}/stderr.txt "
            f"&"))
        print("", file=sys.stderr)

        return SSHDetachedJob(job_id, self, client, work_dir_path)

    def _restore_detached_job(self, job_id: str, job_essence: JsonValue) -> SSHDetachedJob | StoppedJob:
        ssh_config = parse_json_value(job_essence['ssh'], SSHConfig)
        work_dir_path = job_essence['work_dir_path']

        try:
            ssh_client = ssh_connect(ssh_config, retry_times=0)
        except (TimeoutError, NoValidConnectionsError):
            return StoppedJob(job_id)

        return SSHDetachedJob(job_id, self, ssh_client, work_dir_path)
